package protocol

import (
	"log"
	"net"
	"sync"
)

type Worker struct {
	managers []string
	conns    map[string]net.Conn
	mu       sync.Mutex
	jobs     map[string]*Job
	jobFunc  func([]byte) ([]byte, error)
}

type Job struct {
	ID      string
	Status  string
	Payload []byte
}

func NewWorker(managers []string, jobFunc func([]byte) ([]byte, error)) *Worker {
	return &Worker{
		managers: managers,
		conns:    make(map[string]net.Conn),
		jobs:     make(map[string]*Job),
		jobFunc:  jobFunc,
	}
}

func (w *Worker) Connect() error {
	for _, manager := range w.managers {
		log.Printf("Connecting to manager at %s", manager)
		conn, err := net.Dial("tcp", manager)
		if err != nil {
			return err
		}
		log.Printf("Connected to manager at %s", manager)
		w.conns[manager] = conn
	}
	return nil
}

func (w *Worker) SendReady() error {
	msg := &Message{
		Command: READY,
		Headers: make(map[string]string),
	}
	for _, conn := range w.conns {
		if err := msg.Encode(conn); err != nil {
			return err
		}
	}
	return nil
}

func (w *Worker) HandleMessages() error {
	for _, conn := range w.conns {
		go w.handleConnection(conn)
	}
	return nil
}

func (w *Worker) handleConnection(conn net.Conn) {
	log.Printf("Handling messages from manager %s", conn.RemoteAddr())
	for {
		msg := &Message{}
		if err := msg.Decode(conn); err != nil {
			log.Printf("Error decoding message: %v", err)
			continue
		}
		log.Printf("Received message from %s: %v", conn.RemoteAddr(), msg.Command)

		switch msg.Command {
		case WAKE:
			w.handleWake(conn)
		case TAKE:
			w.handleTake(conn, msg)
		case NO:
			log.Println("Received NO command")
		case STATUS:
			w.handleStatus(conn, msg)
		case CANCEL:
			w.handleCancel(conn, msg)
		default:
			log.Printf("Unknown command received: %d", msg.Command)
		}
	}
}

func (w *Worker) handleStatus(conn net.Conn, msg *Message) {
	jobID := msg.Headers["job_id"]
	w.mu.Lock()
	job, exists := w.jobs[jobID]
	w.mu.Unlock()

	statusMsg := &Message{
		Command: STATUS,
		Headers: map[string]string{"job_id": jobID},
	}

	if exists {
		statusMsg.Headers["status"] = job.Status
	} else {
		statusMsg.Headers["status"] = "unknown"
	}

	if err := statusMsg.Encode(conn); err != nil {
		log.Printf("Error sending status message: %v", err)
	}
}

func (w *Worker) handleCancel(conn net.Conn, msg *Message) {
	jobID := msg.Headers["job_id"]
	w.mu.Lock()
	delete(w.jobs, jobID)
	w.mu.Unlock()

	cancelMsg := &Message{
		Command: CANCEL,
		Headers: map[string]string{"job_id": jobID, "status": "cancelled"},
	}

	if err := cancelMsg.Encode(conn); err != nil {
		log.Printf("Error sending cancel message: %v", err)
	}
}

func (w *Worker) handleWake(conn net.Conn) {
	// Check if worker can do the job
	// If yes, send GIVE command
	giveMsg := &Message{
		Command: GIVE,
		Headers: make(map[string]string),
	}
	giveMsg.Encode(conn)
}

func (w *Worker) handleTake(conn net.Conn, msg *Message) {
	jobID := msg.Headers["job_id"]
	log.Printf("Received job %s from manager %s", jobID, conn.RemoteAddr())
	result, err := w.jobFunc(msg.Payload)

	var status string
	if err != nil {
		status = "error"
		log.Printf("Error processing job %s: %v", jobID, err)
	} else {
		status = "completed"
		log.Printf("Job %s completed successfully", jobID)
	}

	doneMsg := &Message{
		Command: DONE,
		Headers: map[string]string{
			"job_id": jobID,
			"status": status,
		},
		Payload: result,
	}

	if err := doneMsg.Encode(conn); err != nil {
		log.Printf("Error sending DONE message: %v", err)
	} else {
		log.Printf("Sent DONE message for job %s to %s", jobID, conn.RemoteAddr())
	}

	// Send READY for the next job
	w.SendReady()
}

func (w *Worker) SendMessage(msg *Message) error {
	w.mu.Lock()
	defer w.mu.Unlock()

	for _, conn := range w.conns {
		if err := msg.Encode(conn); err != nil {
			return err
		}
	}
	return nil
}

func (w *Worker) Close() {
	for _, conn := range w.conns {
		conn.Close()
	}
}
